###
# FROM URL: [python3 gmssl 国密sm4加密与解密]https://zhuanlan.zhihu.com/p/449862473
###

import base64
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import json

# 建议环境变量
SM_KEY = b"""[P4]l,H$U2Of(xn<X-uRBL#sX8Cd'n-{M.2l%z-offxkuQ!e})w^>1q49/MO<yQesWE:hCz?xLg$A$-fWQib7Zx!lWc76bptUwK*}_LR10IY7a1O-k#Y]d,RZT+oI-j2"""


def encode_sm4(value: str) -> str:
    value_bytes = json.dumps(value).encode("utf-8")
    crypt_sm4 = CryptSM4()
    crypt_sm4.set_key(SM_KEY, SM4_ENCRYPT)
    encrypted_bytes = crypt_sm4.crypt_ecb(value_bytes)
    return base64.b64encode(encrypted_bytes).decode("utf-8")


def decode_sm4(encoded: str) -> str:
    encrypted_bytes = base64.b64decode(encoded)
    crypt_sm4 = CryptSM4()
    crypt_sm4.set_key(SM_KEY, SM4_DECRYPT)
    decrypted_bytes = crypt_sm4.crypt_ecb(encrypted_bytes)
    return json.loads(decrypted_bytes.decode("utf-8"))


if __name__ == "__main__":
    text = "123456"

    encrypted = encode_sm4(text)
    print("加密结果：", encrypted)

    decrypted = decode_sm4(encrypted)
    print("解密结果：", decrypted == text)
